from functools import wraps
from flask import Flask, request, jsonify
import flask
import json
import requests
from typing import Optional
import warnings
from cache import MemoryCache
from flask_cors import CORS
from metrics import MetricsService
import pandas as pd
import plotly
import plotly.express as px
import plotly.graph_objects as go
from langflow_service import LangflowService
from datetime import datetime
from visualization_service import VisualizationService
from config import ChromaConfig
from chroma_manager import ChromaManager

app = Flask(__name__, static_url_path='')
CORS(app)

# SETUP
cache = MemoryCache()

# Initialize service
langflow_service = LangflowService("http://localhost:7860")

# Initialize services
chroma_config = ChromaConfig(
    host=None,
    ollama_base_url="http://localhost:11434",
    embedding_model="viosay/conan-embedding-v1"
)
chroma = ChromaManager(chroma_config)

viz_service = VisualizationService(
    chroma_manager=chroma,
    model_service_type="ollama",
    base_url="http://localhost:11434",
    model_name="qwen2.5:3b"
)

# 原来的 run_flow 函数
BASE_API_URL = "http://81.70.215.246:40006"
FLOW_ID = "563537f3-522b-43ec-8b34-3c10f7743fcd"
ENDPOINT = ""  # You can set a specific endpoint name in the flow settings

# You can tweak the flow by adding a tweaks dictionary
TWEAKS = {
    "ChatInput-SFFoz": {},
    "ChatOutput-cTdDv": {},
    "Prompt-FKbfn": {}
}

def run_flow(message: str,
             endpoint: str,
             output_type: str = "chat",
             input_type: str = "chat",
             tweaks: Optional[dict] = None,
             api_key: Optional[str] = None) -> dict:
    """
    Run a flow with a given message and optional tweaks.

    :param message: The message to send to the flow
    :param endpoint: The ID or the endpoint name of the flow
    :param tweaks: Optional tweaks to customize the flow
    :return: The JSON response from the flow
    """
    api_url = f"{BASE_API_URL}/api/v1/run/{endpoint}"

    payload = {
        "input_value": message,
        "output_type": output_type,
        "input_type": input_type,
    }
    headers = None
    if tweaks:
        payload["tweaks"] = tweaks
    if api_key:
        headers = {"x-api-key": api_key}
    response = requests.post(api_url, json=payload, headers=headers)
    return response.json()

# 创建 Flask 实例
app = Flask(__name__)

# 设置一个 POST 接口来接收消息并返回输出
@app.route('/run_flow', methods=['POST'])
def run_flow_service():
    # 获取请求体的 JSON 数据
    data = request.get_json()

    # 从请求体中提取参数
    message = data.get('message')
    endpoint = data.get('endpoint', FLOW_ID)
    output_type = data.get('output_type', 'chat')
    input_type = data.get('input_type', 'chat')
    tweaks = data.get('tweaks', TWEAKS)
    api_key = data.get('api_key', None)

    # 调用原始的 run_flow 函数
    response = run_flow(
        message=message,
        endpoint=endpoint,
        output_type=output_type,
        input_type=input_type,
        tweaks=tweaks,
        api_key=api_key
    )

    # 提取 output 数据
    output = response['outputs'][0]['outputs'][0]['results']['message']['data']['text']

    # 返回响应数据
    return jsonify({"output": output})

def requires_cache(fields):
    def decorator(f):
        @wraps(f)
        def decorated(*args, **kwargs):
            id = request.args.get('id')

            if id is None:
                return jsonify({"type": "error", "error": "No id provided"})
            
            for field in fields:
                if cache.get(id=id, field=field) is None:
                    return jsonify({"type": "error", "error": f"No {field} found"})
            
            field_values = {field: cache.get(id=id, field=field) for field in fields}
            
            # Add the id to the field_values
            field_values['id'] = id

            return f(*args, **field_values, **kwargs)
        return decorated
    return decorator

@app.route('/api/v0/generate_params', methods=['GET'])
async def generate_params():
    question = flask.request.args.get('question')

    if question is None:
        return jsonify({"type": "error", "error": "No question provided"})

    id = cache.generate_id(question=question)

    # params = {
    #         "startDate": '2022-01-01',
    #         "endDate": '2025-01-01',
    #         "kpis": 'S0010001 as kpi1'
    #     }

    # Get query parameters
    params = viz_service.generate_kpi_params(question)

    cache.set(id=id, field='question', value=question)
    cache.set(id=id, field='params', value=params)

    return jsonify(
        {
            "type": "params", 
            "id": id,
            "text": params,
        })

@app.route('/api/v0/get_kpis', methods=['GET'])
@requires_cache(['params'])
def get_kpis(id: str, params):
    try:
        client = MetricsService()

        # Get KPI data
        df = client.get_kpis(params)

        cache.set(id=id, field='df', value=df)

        return jsonify(
            {
                "type": "df", 
                "id": id,
                "df": df.head(10).to_json(orient='records'),
            })

    except Exception as e:
        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/generate_plotly_figure', methods=['GET'])
@requires_cache(['df', 'question', 'params'])
async def generate_plotly_figure(id: str, df, question, params):
    try:
        # code = vn.generate_plotly_code(question=question, sql=sql, df_metadata=f"Running df.dtypes gives:\n {df.dtypes}")
        # code = "```python\nimport plotly.express as px\nimport plotly.graph_objects as go\n\nif len(df) == 1:\n    fig = go.Figure(go.Indicator(\n        mode = \"number\",\n        value = df['visit_count'].iloc[0],\n        title = {\"text\": \"Daily Visit Count<br><span style='font-size:0.8em;color:gray'>Date: \" + str(df['date'].iloc[0]) + \"</span>\"},\n        domain = {'x': [0, 1], 'y': [0, 1]}\n    ))\nelse:\n    fig = px.line(df, x='date', y='visit_count', title='Daily Visit Count')\n    fig.update_xaxes(title_text='Date')\n    fig.update_yaxes(title_text='Visit Count')\n\nfig.show()\n```"
        # code = await langflow_service.generate_plot_code(
        #     question,
        #     kpi_params,
        #     df_metadata
        # )
        code = viz_service.generate_plot_code(
                question=question,
                query_params=params,
                df_metadata=f"Running df.dtypes gives:\n {df.dtypes}"
            )

        fig = get_plotly_figure(plotly_code=code, df=df, dark_mode=False)
        fig_json = fig.to_json()

        cache.set(id=id, field='fig_json', value=fig_json)

        return jsonify(
            {
                "type": "plotly_figure", 
                "id": id,
                "fig": fig_json,
            })
    except Exception as e:
        # Print the stack trace
        import traceback
        traceback.print_exc()

        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/run_all', methods=['GET'])
def run_all():
    try:
        # Step 1: Call generate_params
        question = flask.request.args.get('question')
        if question is None:
            return jsonify({"type": "error", "error": "No question provided"})

        generate_kpis_response = generate_params()
        if generate_kpis_response.status_code != 200:
            return jsonify({"type": "error", "error": "Failed to generate KPIs"})

        generate_kpis_data = generate_kpis_response.get_json()
        if generate_kpis_data.get('type') == 'error':
            return jsonify(generate_kpis_data)

        id = generate_kpis_data.get('id')

        # Step 2: Call get_kpis
        with app.test_request_context(f'/api/v0/get_kpis?id={id}', method='GET'):
            get_kpis_response = get_kpis()

        if get_kpis_response.status_code != 200:
            return jsonify({"type": "error", "error": "Failed to retrieve KPIs"})

        get_kpis_data = get_kpis_response.get_json()
        if get_kpis_data.get('type') == 'error':
            return jsonify(get_kpis_data)

        # Step 3: Call generate_plotly_figure and directly return its response
        with app.test_request_context(f'/api/v0/generate_plotly_figure?id={id}', method='GET'):
            generate_plotly_figure_response = generate_plotly_figure()

        if generate_plotly_figure_response.status_code != 200:
            return jsonify({"type": "error", "error": "Failed to generate Plotly figure"})

        return generate_plotly_figure_response  # Directly return the response from the last function

    except Exception as e:
        return jsonify({"type": "error", "error": str(e)})
    
def get_plotly_figure(
        plotly_code: str, df: pd.DataFrame, dark_mode: bool = True
    ) -> plotly.graph_objs.Figure:
        """
        **Example:**
        ```python
        fig = vn.get_plotly_figure(
            plotly_code="fig = px.bar(df, x='name', y='salary')",
            df=df
        )
        fig.show()
        ```
        Get a Plotly figure from a dataframe and Plotly code.

        Args:
            df (pd.DataFrame): The dataframe to use.
            plotly_code (str): The Plotly code to use.

        Returns:
            plotly.graph_objs.Figure: The Plotly figure.
        """
        ldict = {"df": df, "px": px, "go": go}
        try:
            exec(plotly_code, globals(), ldict)

            fig = ldict.get("fig", None)
        except Exception as e:
            # Inspect data types
            numeric_cols = df.select_dtypes(include=["number"]).columns.tolist()
            categorical_cols = df.select_dtypes(
                include=["object", "category"]
            ).columns.tolist()

            # Decision-making for plot type
            if len(numeric_cols) >= 2:
                # Use the first two numeric columns for a scatter plot
                fig = px.scatter(df, x=numeric_cols[0], y=numeric_cols[1])
            elif len(numeric_cols) == 1 and len(categorical_cols) >= 1:
                # Use a bar plot if there's one numeric and one categorical column
                fig = px.bar(df, x=categorical_cols[0], y=numeric_cols[0])
            elif len(categorical_cols) >= 1 and df[categorical_cols[0]].nunique() < 10:
                # Use a pie chart for categorical data with fewer unique values
                fig = px.pie(df, names=categorical_cols[0])
            else:
                # Default to a simple line plot if above conditions are not met
                fig = px.line(df)

        if fig is None:
            return None

        if dark_mode:
            fig.update_layout(template="plotly_dark")

        return fig


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)
